package com.github.defense.plugin.gray.netflix.ribbon;

import com.github.defense.plugin.gray.choose.GrayServerChoose;
import com.github.defense.plugin.gray.constants.GrayClientConstant;
import com.github.defense.plugin.gray.holder.GrayClientHolder;
import com.github.defense.plugin.gray.holder.GrayEnvContextHolder;
import com.google.common.base.Optional;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;


/**
 * 灰度路由决策
 *
 * @Author: lettger
 * @Date: 2021/6/5 2:14 下午
 */
public class GrayRibbonRule extends ZoneAvoidanceRule {

    private GrayServerChoose<Server> choose;

    public GrayRibbonRule() {
        this.choose = GrayClientHolder.getServerChoose();
    }

    @Override
    public Server choose(final Object key) {
        // 未开启灰度
        if (!GrayEnvContextHolder.getEnvBool(GrayClientConstant.GRAY_CLIENT_ENABLE)) {
            return super.choose(key);
        }
        return choose.choose(getLoadBalancer().getReachableServers(),
                (servers) -> {
                    Optional<Server> server = getPredicate().chooseRoundRobinAfterFiltering(servers, key);
                    if (server.isPresent()) {
                        return server.get();
                    }
                    return null;
                }
        );
    }
}
